from matplotlib import pyplot as plt

def read_file(dataset, model, noise_type, noise_rate):
    loss = []
    epoch = []
    num = [[], [], [], [], []]
    data = [[], [], [], [], []]
    
    path = '../results/'+dataset+'/'+model+'/'+dataset+'_'+model+'_'+noise_type+'_'+noise_rate+'.txt'
    with open(path, 'r') as f:
        '''
        epoch: clean_true clean_false noise_true noise_false noise_other train_acc train_loss test_acc test_loss num1 num2 num3 num4 num5
        1: -0.5490396022796631 0.9847707748413086 0.4409019351005554 -0.4204067587852478 0.8046616911888123 21.878000259399414 2.3153586387634277 29.790000915527344 0.0003822564030997455 4288.0 35674.0 1067.0 1026.0 7945.0
        '''
        lines = f.readlines()
        for line in lines[1:]:
            tmp = line.strip().split()
            epoch.append(int(tmp[0][:-1]))
            for i in range(5):
                data[i].append(float(tmp[i + 1]))
                num[i].append(float(tmp[i - 5]))
            loss.append(float(tmp[-6]))
    title = dataset+'_'+model+'_'+noise_type+'_'+noise_rate
    plt.figure(figsize=(30, 10))
    plt.title(title)
    ax_list = [plt.subplot(231), plt.subplot(232), plt.subplot(234), plt.subplot(235), plt.subplot(236), plt.subplot(233)]
    color = ['y', 'r', 'g', 'b', 'c', 'm', 'k', 'w']
    label = ['clean_true', 'clean_false', 'noise_true', 'noise_false', 'noise_other']

    for i in range(5):
        ax_list[i].set_xlabel('epoch')
        ax_list[i].set_ylabel('cosine sim', color=color[1])
        ax_list[i].plot(epoch, data[i], color[1], linewidth=1, label=label[i]+'_sim')
        tmp = ax_list[i].twinx()
        tmp.set_ylabel('num', color=color[3])
        tmp.plot(epoch, num[i], color=color[3], linewidth=1, label=label[i]+'_num')
        ax_list[i].legend(loc='lower right')
        tmp.legend(loc='upper right')
    ax_list[5].set_xlabel('epoch')
    ax_list[5].set_ylabel('loss')
    ax_list[5].plot(epoch, loss, linewidth=1)
    plt.savefig('5_graphs/'+title+'.png')
    plt.close()

    return data, epoch, loss

read_file('cifar10', 'resnet34', 'symmetric', '0.2')
read_file('cifar100', 'resnet34', 'symmetric', '0.2')
read_file('cifar10', 'resnet50', 'symmetric', '0.2')
read_file('cifar100', 'resnet50', 'symmetric', '0.2')

read_file('cifar10', 'resnet34', 'symmetric', '0.6')
read_file('cifar100', 'resnet34', 'symmetric', '0.6')
read_file('cifar10', 'resnet50', 'symmetric', '0.6')
read_file('cifar100', 'resnet50', 'symmetric', '0.6')

read_file('cifar10', 'resnet34', 'clean', '0.0')
read_file('cifar100', 'resnet34', 'clean', '0.0')
read_file('cifar10', 'resnet50', 'clean', '0.0')
read_file('cifar100', 'resnet50', 'clean', '0.0')